import requests
import json
import os
import base64

API_KEY = "QZ2m1g1wRgX5UoJDfcPTT1Kd"
SECRET_KEY = "8B9bHYFaePqtrOQMEXQHDGvZ5UGAo7zs"


def main():
    progress = 3

    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/text2image/sd_xl?access_token=" + get_access_token()

    payload = json.dumps({
        "prompt": "Unitree humanoid robot sits on a gray sofa with some electrical ornaments in the background",
        "size": "768x768",
        "n": 1,
        "steps": 20,
        "sampler_index": "Euler a"
    }, ensure_ascii=False)
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload.encode("utf-8"))

    # print(response.text)
    json_response = response.text

    # 解析 JSON 响应
    data = json.loads(json_response)

    # 提取 b64_image 数据
    b64_images = [item["b64_image"] for item in data["data"]]

    # 确保保存路径存在
    save_path = "/root/data/sd_recent"
    os.makedirs(save_path, exist_ok=True)

    # 遍历每个 b64_image 并保存为图片
    for i, b64_image in enumerate(b64_images):
        image_path = os.path.join(save_path, f"{progress}.png")

        try:
            # 将 base64 数据解码为二进制数据
            image_data = base64.b64decode(b64_image)

            # 保存图片
            with open(image_path, "wb") as f:
                f.write(image_data)
            print(f"图片已保存到: {image_path}")
        except Exception as e:
            print(f"保存图片时出错: {e}")


def get_access_token():
    """
    使用 AK，SK 生成鉴权签名（Access Token）
    :return: access_token，或是None(如果错误)
    """
    url = "https://aip.baidubce.com/oauth/2.0/token"
    params = {"grant_type": "client_credentials", "client_id": API_KEY, "client_secret": SECRET_KEY}
    return str(requests.post(url, params=params).json().get("access_token"))


if __name__ == '__main__':
    main()
